// 通过缩减Tensor ShapeInfo维度，优化栈空间
#define K_MAX_SHAPE_DIM 0

#include "kernel_operator.h"

using namespace AscendC;

#define ll long long

constexpr int32_t BUFFER_NUM = 1;
constexpr int32_t TQUE_NUM = 1;
constexpr int32_t BUFFER_NUM_SP = 2;


template<typename T> class CopySign {
private:
    TPipe* pipe;
    TQue<QuePosition::VECIN, TQUE_NUM> inQueueX1;
    TQue<QuePosition::VECIN, TQUE_NUM> inQueueX2;
    TQue<QuePosition::VECOUT, TQUE_NUM> outQueueY;
    GlobalTensor<T> x1Gm;
    GlobalTensor<T> x2Gm;
    GlobalTensor<T> yGm;
    int blockIdx;
    ll offset;
    ll offset1;
    ll offset2;
    int indices[5];
    ll totalLength1;
    ll totalLength2;
    ll totalLength;
    int calcLength;
    int dimNum1;
    int dimNum2;
    int dimNum;
    int* dim1;
    int* dim2;
    int* dim;

public:
    __aicore__ inline CopySign() {}
    __aicore__ inline void Init(GM_ADDR x1, GM_ADDR x2, GM_ADDR y, GM_ADDR workspace, TPipe* pipeIn, int* dim2, int* dim)
    {
        pipe = pipeIn;
        this->dimNum = 4;
        this->dim2 = dim2;
        this->dim = dim;

        this->totalLength = 1;
        this->totalLength2 = 1;
        for (int i = 0; i < dimNum; i++) {
            this->totalLength *= dim[i];
            this->totalLength2 *= dim2[i];
        }
        // 打印每个dimnum和dim
        // printf("dimNum1: %d, dimNum2: %d, dimNum: %d\n", dimNum1, dimNum2, dimNum);
        // for (int i = 0; i < dimNum; i++) {
        //     printf("dim[%d]: %d, dim1[%d]: %d, dim2[%d]: %d\n", i, dim[i], i, dim1[i], i, dim2[i]);
        // }
        
        blockIdx = GetBlockIdx();
        calcLength = dim[dimNum - 1];
        offset1 = blockIdx * calcLength;
        // 计算indices（每一维的索引）
        indices[dimNum - 1] = 0;
        int tmp = blockIdx;
        for (int i = dimNum - 2; i >= 0; i--) {
            indices[i] = tmp % dim[i];
            tmp /= dim[i];
        }
        // printf("blockIdx: %d, indices: ", blockIdx);
        // for (int i = 0; i < dimNum; i++) {
        //     printf("%d ", indices[i]);
        // }
        // printf("\n");

        // 计算x2索引
        offset2 = 0;
        ll stride2 = 1;
        for (int i = dimNum - 1; i >= 0; i--) {
            int idx = indices[i] % dim2[i]; // 处理广播
            offset2 += idx * stride2;
            stride2 *= dim2[i];
        }
        // printf("offset1: %d, offset2: %d\n", offset1, offset2);
        

        x1Gm.SetGlobalBuffer((__gm__ T *)x1 + offset1, calcLength * sizeof(T));
        x2Gm.SetGlobalBuffer((__gm__ T *)x2 + offset2, dim2[dimNum - 1] * sizeof(T));
        yGm.SetGlobalBuffer((__gm__ T *)y + offset1, calcLength * sizeof(T));

        pipe->InitBuffer(inQueueX1, BUFFER_NUM, calcLength * sizeof(T));
        pipe->InitBuffer(inQueueX2, BUFFER_NUM, calcLength * sizeof(T));
        pipe->InitBuffer(outQueueY, BUFFER_NUM, calcLength * sizeof(T));
    }
    __aicore__ inline void Process()
    {
        DataCopyPadExtParams<T> padParams{false, 0, 0, 0};
        DataCopyExtParams copyParams = {
            (uint16_t)1, 
            (uint32_t)(calcLength * sizeof(T)), 
            0,
            0,
            0
        };
        DataCopyExtParams copyParams2 = {
            (uint16_t)1, 
            (uint32_t)(1 * sizeof(T)), 
            0,
            0,
            0
        };
        if (dim2[dimNum - 1] == 1) {
            // 处理x2为1的情况
            LocalTensor<T> x1Local = inQueueX1.AllocTensor<T>();
            LocalTensor<T> x2Local = inQueueX2.AllocTensor<T>();
            DataCopyPad(x1Local, x1Gm, copyParams, padParams);
            inQueueX1.EnQue(x1Local);
            x1Local = inQueueX1.DeQue<T>();
            DataCopyPad(x2Local, x2Gm, copyParams2, padParams);
            inQueueX2.EnQue(x2Local);
            x2Local = inQueueX2.DeQue<T>();
            T x2 = x2Local(0);
            inQueueX2.FreeTensor(x2Local);
            LocalTensor<T> yLocal = outQueueY.AllocTensor<T>();

            bool x2Sign;
            if constexpr (Std::is_same<T, bfloat16_t>::value) {
                float x2f = ToFloat(x2);
                x2Sign = (x2f >= 0.0f);
            } else if constexpr (Std::is_same<T, float>::value) {
                x2Sign = (x2 >= 0.0f);
            } else {
                x2Sign = (float)x2 >= 0.0f;
            }

            for (int i = 0; i < calcLength; i++) {
                T x1 = x1Local(i);
                if constexpr (Std::is_same<T, bfloat16_t>::value) {
                    float x1f = ToFloat(x1);
                    if (x2Sign) {
                        yLocal(i) = ToBfloat16(x1f >= 0.0f ? x1f : -x1f);
                    } else {
                        yLocal(i) = ToBfloat16(x1f >= 0.0f ? -x1f : x1f);
                    }
                } else if constexpr (Std::is_same<T, float>::value) {
                    yLocal(i) = x2Sign ? (x1 >= 0.0f ? x1 : -x1) : (x1 >= 0.0f ? -x1 : x1);
                } else {
                    float x1f = (float)x1;
                    float x2f = (float)x2;
                    yLocal(i) = (half)(x2f >= 0.0f ? (x1f >= 0.0f ? x1f : -x1f) : (x1f >= 0.0f ? -x1f : x1f));
                }
            }

            inQueueX1.FreeTensor(x1Local);
            outQueueY.EnQue(yLocal);
            yLocal = outQueueY.DeQue<T>();
            DataCopyPad(yGm, yLocal, copyParams);
            outQueueY.FreeTensor(yLocal);
            return;
        }
        // 处理x2为多维的情况
        LocalTensor<T> x1Local = inQueueX1.AllocTensor<T>();
        LocalTensor<T> x2Local = inQueueX2.AllocTensor<T>();
        DataCopyPad(x1Local, x1Gm, copyParams, padParams);
        DataCopyPad(x2Local, x2Gm, copyParams, padParams);
        inQueueX1.EnQue(x1Local);
        inQueueX2.EnQue(x2Local);
        x1Local = inQueueX1.DeQue<T>();
        x2Local = inQueueX2.DeQue<T>();
        LocalTensor<T> yLocal = outQueueY.AllocTensor<T>();

        for (int i = 0; i < calcLength; i++) {
            T x1 = x1Local(i);
            T x2 = x2Local(i);
            if constexpr (Std::is_same<T, bfloat16_t>::value) {
                float x1f = ToFloat(x1);
                float x2f = ToFloat(x2);
                yLocal(i) = x2f >= 0.0f ? ToBfloat16(x1f >= 0.0f ? x1f : -x1f) : ToBfloat16(x1f >= 0.0f ? -x1f : x1f);
            } else if constexpr (Std::is_same<T, float>::value) {
                yLocal(i) = x2 >= 0.0f ? (x1 >= 0.0f ? x1 : -x1) : (x1 >= 0.0f ? -x1 : x1);
            } else {
                float x1f = (float)x1;
                float x2f = (float)x2;
                yLocal(i) = (half)(x2f >= 0.0f ? (x1f >= 0.0f ? x1f : -x1f) : (x1f >= 0.0f ? -x1f : x1f));
            }
        }
        
        outQueueY.EnQue(yLocal);
        inQueueX1.FreeTensor(x1Local);
        inQueueX2.FreeTensor(x2Local);
        yLocal = outQueueY.DeQue<T>();
        DataCopyPad(yGm, yLocal, copyParams);
        outQueueY.FreeTensor(yLocal);
    }
};


template<typename T> class CopySign_SP {
private:
    TPipe* pipe;
    TQue<QuePosition::VECIN, TQUE_NUM> inQueueX1;
    TQue<QuePosition::VECIN, TQUE_NUM> inQueueX2;
    TQue<QuePosition::VECOUT, TQUE_NUM> outQueueY;
    GlobalTensor<T> x1Gm;
    GlobalTensor<T> x2Gm;
    GlobalTensor<T> yGm;
    int blockNum;
    int blockIdx;
    ll offset;
    ll offset1;
    ll offset2;
    int indices[5];
    ll totalLength1;
    ll totalLength2;
    ll totalLength;
    int tileLength;
    int dimNum1;
    int dimNum2;
    int dimNum;
    int* dim1;
    int* dim2;
    int* dim;

public:
    __aicore__ inline CopySign_SP() {}
    __aicore__ inline void Init(GM_ADDR x1, GM_ADDR x2, GM_ADDR y, GM_ADDR workspace, TPipe* pipeIn, int* dim2, int* dim)
    {
        pipe = pipeIn;
        this->dimNum = 4;
        this->dim2 = dim2;
        this->dim = dim;

        this->totalLength = 1;
        this->totalLength2 = 1;
        for (int i = 0; i < dimNum; i++) {
            this->totalLength *= dim[i];
            this->totalLength2 *= dim2[i];
        }
        // 打印每个dimnum和dim
        // printf("dimNum1: %d, dimNum2: %d, dimNum: %d\n", dimNum1, dimNum2, dimNum);
        // for (int i = 0; i < dimNum; i++) {
        //     printf("dim[%d]: %d, dim1[%d]: %d, dim2[%d]: %d\n", i, dim[i], i, dim1[i], i, dim2[i]);
        // }

        blockNum = GetBlockNum();
        blockIdx = GetBlockIdx();
        tileLength = 4096;
        if (tileLength > totalLength) {
            tileLength = totalLength;
        }

        x1Gm.SetGlobalBuffer((__gm__ T *)x1, totalLength);
        x2Gm.SetGlobalBuffer((__gm__ T *)x2, totalLength);
        yGm.SetGlobalBuffer((__gm__ T *)y, totalLength);

        pipe->InitBuffer(inQueueX1, BUFFER_NUM_SP, tileLength * sizeof(T));
        pipe->InitBuffer(inQueueX2, BUFFER_NUM_SP, tileLength * sizeof(T));
        pipe->InitBuffer(outQueueY, BUFFER_NUM_SP, tileLength * sizeof(T));
    }
    __aicore__ inline void Process()
    {
        #define min(x, y) ((x) < (y) ? (x) : (y))
        for (size_t offset = blockIdx * tileLength; offset < totalLength; offset += blockNum * tileLength) {
            int calcLength = min(tileLength, totalLength - offset);
            Compute(offset, calcLength);
        }
    }
    __aicore__ inline void Compute(size_t offset, int calcLength)
    {
        DataCopyPadExtParams<T> padParams{false, 0, 0, 0};
        DataCopyExtParams copyParams = {
            (uint16_t)1, 
            (uint32_t)(calcLength * sizeof(T)), 
            0,
            0,
            0
        };
        // 处理x2为多维的情况
        LocalTensor<T> x1Local = inQueueX1.AllocTensor<T>();
        DataCopyPad(x1Local, x1Gm[offset], copyParams, padParams);
        inQueueX1.EnQue(x1Local);
        LocalTensor<T> x2Local = inQueueX2.AllocTensor<T>();
        DataCopyPad(x2Local, x2Gm[offset], copyParams, padParams);
        inQueueX2.EnQue(x2Local);
        LocalTensor<T> yLocal;

        if constexpr (std::is_same_v<T, float>) {
            x1Local = inQueueX1.DeQue<T>();
            ShiftLeft(x1Local.template ReinterpretCast<uint32_t>(), x1Local.template ReinterpretCast<uint32_t>(), (uint32_t)1, calcLength);
            ShiftRight(x1Local.template ReinterpretCast<uint32_t>(), x1Local.template ReinterpretCast<uint32_t>(), (uint32_t)1, calcLength);
            x2Local = inQueueX2.DeQue<T>();
            ShiftRight(x2Local.template ReinterpretCast<uint32_t>(), x2Local.template ReinterpretCast<uint32_t>(), (uint32_t)31, calcLength);
            ShiftLeft(x2Local.template ReinterpretCast<uint32_t>(), x2Local.template ReinterpretCast<uint32_t>(), (uint32_t)31, calcLength);
            yLocal = outQueueY.AllocTensor<T>();
            Or(yLocal.template ReinterpretCast<uint16_t>(), x1Local.template ReinterpretCast<uint16_t>(), x2Local.template ReinterpretCast<uint16_t>(), calcLength * 2);
        } else {
            x1Local = inQueueX1.DeQue<T>();
            ShiftLeft(x1Local.template ReinterpretCast<uint16_t>(), x1Local.template ReinterpretCast<uint16_t>(), (uint16_t)1, calcLength);
            ShiftRight(x1Local.template ReinterpretCast<uint16_t>(), x1Local.template ReinterpretCast<uint16_t>(), (uint16_t)1, calcLength);
            x2Local = inQueueX2.DeQue<T>();
            ShiftRight(x2Local.template ReinterpretCast<uint16_t>(), x2Local.template ReinterpretCast<uint16_t>(), (uint16_t)15, calcLength);
            ShiftLeft(x2Local.template ReinterpretCast<uint16_t>(), x2Local.template ReinterpretCast<uint16_t>(), (uint16_t)15, calcLength);
            yLocal = outQueueY.AllocTensor<T>();
            Or(yLocal.template ReinterpretCast<uint16_t>(), x1Local.template ReinterpretCast<uint16_t>(), x2Local.template ReinterpretCast<uint16_t>(), calcLength);
        }

        inQueueX1.FreeTensor(x1Local);
        inQueueX2.FreeTensor(x2Local);
        outQueueY.EnQue(yLocal);
        yLocal = outQueueY.DeQue<T>();
        DataCopyPad(yGm[offset], yLocal, copyParams);
        outQueueY.FreeTensor(yLocal);
    }
};


template<typename T> class CopySign_CS3 {
private:
    TPipe* pipe;
    TQue<TPosition::VECIN, TQUE_NUM> inQueueX1, inQueueX2;
    TQue<TPosition::VECOUT, TQUE_NUM> outQueueY;
    TQue<TPosition::VECCALC, TQUE_NUM> calcQueueX1, calcQueueX2, calcQueueY;
    GlobalTensor<T> x1Gm, x2Gm, yGm;
    int indices[5];
    ll totalLength1;
    ll totalLength2;
    ll totalLength;
    int dimNum1;
    int dimNum2;
    int dimNum;
    int* dim1;
    int* dim2;
    int* dim;
    int blockNum;
    int blockIdx;
    int blockLength;
    int tailBlockNum;
    int tileLength;
    int tileNum;
    int calcLength;
    int bcLength;
    int lastLength;
    int beforeBC;
    int lastLengthAlign;
    int blockLengthAlign;
    int tileLengthAlign;
    int idx0;
    int idx1;

public:
    __aicore__ inline CopySign_CS3() {}
    __aicore__ inline void Init(GM_ADDR x1, GM_ADDR x2, GM_ADDR y, GM_ADDR workspace, TPipe* pipeIn, int* dim2, int* dim)
    {
        pipe = pipeIn;
        this->dimNum = 4;
        this->dim2 = dim2;
        this->dim = dim;

        this->totalLength = 1;
        this->totalLength2 = 1;
        for (int i = 0; i < dimNum; i++) {
            this->totalLength *= dim[i];
            this->totalLength2 *= dim2[i];
        }

        blockNum = GetBlockNum();
        blockIdx = GetBlockIdx();
        calcLength = dim[2];
        lastLength = dim[2];
        tileLength = lastLength;
        bcLength = dim[1];
        beforeBC = dim[0];
        // this->blockLength = blockLength;
        blockLength = (beforeBC * bcLength + blockNum - 1) / blockNum;
        int alreadyLength = blockIdx * blockLength;
        idx0 = alreadyLength / bcLength;
        idx1 = alreadyLength % bcLength;

        // printf("blockIdx = %d, blockLength = %d, blockLengthAlign = %d\n", blockIdx, blockLength, blockLengthAlign);
        // printf("alreadyLength = %d, tileLength = %d, tileLengthAlign = %d\n", alreadyLength, tileLength, tileLengthAlign);

        x1Gm.SetGlobalBuffer((__gm__ T *)x1, totalLength);
        x2Gm.SetGlobalBuffer((__gm__ T *)x2, totalLength2);
        yGm.SetGlobalBuffer((__gm__ T *)y, totalLength);

        pipe->InitBuffer(inQueueX1, BUFFER_NUM_SP, tileLength * sizeof(T));
        pipe->InitBuffer(inQueueX2, BUFFER_NUM_SP, tileLength * sizeof(T));
        pipe->InitBuffer(outQueueY, BUFFER_NUM_SP, tileLength * sizeof(T));
    }
    __aicore__ inline void Process()
    {
        DataCopyPadExtParams<T> padParams{false, 0, 0, 0};
        DataCopyExtParams copyParams = {
            (uint16_t)1, 
            (uint32_t)(lastLength * sizeof(T)), 
            0,
            0,
            0
        };
        LocalTensor<T> x2Local = inQueueX2.AllocTensor<T>();
        DataCopyPad(x2Local, x2Gm[idx0 * lastLength], copyParams, padParams);
        inQueueX2.EnQue(x2Local);
        x2Local = inQueueX2.DeQue<T>();
        if constexpr (std::is_same_v<T, float>) {
            ShiftRight(x2Local.template ReinterpretCast<uint32_t>(), x2Local.template ReinterpretCast<uint32_t>(), (uint32_t)31, calcLength);
            ShiftLeft(x2Local.template ReinterpretCast<uint32_t>(), x2Local.template ReinterpretCast<uint32_t>(), (uint32_t)31, calcLength);
        } else {
            ShiftRight(x2Local.template ReinterpretCast<uint16_t>(), x2Local.template ReinterpretCast<uint16_t>(), (uint16_t)15, calcLength);
            ShiftLeft(x2Local.template ReinterpretCast<uint16_t>(), x2Local.template ReinterpretCast<uint16_t>(), (uint16_t)15, calcLength);
        }

        int offset1 = idx0 * bcLength * lastLength + idx1 * lastLength;
        for (int i = 0; i < blockLength && offset1 < totalLength; i++, offset1 += lastLength) {
            LocalTensor<T> yLocal;
            LocalTensor<T> x1Local = inQueueX1.AllocTensor<T>();
            DataCopyPad(x1Local, x1Gm[offset1], copyParams, padParams);
            inQueueX1.EnQue(x1Local);
            x1Local = inQueueX1.DeQue<T>();
            if constexpr (std::is_same_v<T, float>) {
                x1Local = inQueueX1.DeQue<T>();
                ShiftLeft(x1Local.template ReinterpretCast<uint32_t>(), x1Local.template ReinterpretCast<uint32_t>(), (uint32_t)1, calcLength);
                ShiftRight(x1Local.template ReinterpretCast<uint32_t>(), x1Local.template ReinterpretCast<uint32_t>(), (uint32_t)1, calcLength);
                yLocal = outQueueY.AllocTensor<T>();
                Or(yLocal.template ReinterpretCast<uint16_t>(), x1Local.template ReinterpretCast<uint16_t>(), x2Local.template ReinterpretCast<uint16_t>(), calcLength * 2);
            } else {
                x1Local = inQueueX1.DeQue<T>();
                ShiftLeft(x1Local.template ReinterpretCast<uint16_t>(), x1Local.template ReinterpretCast<uint16_t>(), (uint16_t)1, calcLength);
                ShiftRight(x1Local.template ReinterpretCast<uint16_t>(), x1Local.template ReinterpretCast<uint16_t>(), (uint16_t)1, calcLength);
                yLocal = outQueueY.AllocTensor<T>();
                Or(yLocal.template ReinterpretCast<uint16_t>(), x1Local.template ReinterpretCast<uint16_t>(), x2Local.template ReinterpretCast<uint16_t>(), calcLength);
            }
            inQueueX1.FreeTensor(x1Local);
            outQueueY.EnQue(yLocal);
            yLocal = outQueueY.DeQue<T>();
            DataCopyPad(yGm[offset1], yLocal, copyParams);
            outQueueY.FreeTensor(yLocal);

            idx1++;
            if (idx1 == bcLength) {
                idx0++;
                idx1 = 0;

                inQueueX2.FreeTensor(x2Local);
                x2Local = inQueueX2.AllocTensor<T>();
                DataCopyPad(x2Local, x2Gm[idx0 * lastLength], copyParams, padParams);
                inQueueX2.EnQue(x2Local);
                x2Local = inQueueX2.DeQue<T>();
                if constexpr (std::is_same_v<T, float>) {
                    ShiftRight(x2Local.template ReinterpretCast<uint32_t>(), x2Local.template ReinterpretCast<uint32_t>(), (uint32_t)31, calcLength);
                    ShiftLeft(x2Local.template ReinterpretCast<uint32_t>(), x2Local.template ReinterpretCast<uint32_t>(), (uint32_t)31, calcLength);
                } else {
                    ShiftRight(x2Local.template ReinterpretCast<uint16_t>(), x2Local.template ReinterpretCast<uint16_t>(), (uint16_t)15, calcLength);
                    ShiftLeft(x2Local.template ReinterpretCast<uint16_t>(), x2Local.template ReinterpretCast<uint16_t>(), (uint16_t)15, calcLength);
                }
            }
        }
    }
};


extern "C" __global__ __aicore__ void copysign(GM_ADDR input, GM_ADDR other, GM_ADDR out, GM_ADDR workspace, GM_ADDR tiling) {
    GET_TILING_DATA(tiling_data, tiling);
    TPipe pipe;
    if (TILING_KEY_IS(1)) {
        CopySign_SP<DTYPE_INPUT> op;
        op.Init(input, other, out, workspace, &pipe, tiling_data.dim2, tiling_data.dim);
        op.Process();
    } else if (TILING_KEY_IS(2)) {
        CopySign_CS3<DTYPE_INPUT> op;
        op.Init(input, other, out, workspace, &pipe, tiling_data.dim2, tiling_data.dim);
        op.Process();
    } else if (TILING_KEY_IS(3)) {
        CopySign<DTYPE_INPUT> op;
        op.Init(input, other, out, workspace, &pipe, tiling_data.dim2, tiling_data.dim);
        op.Process();
    }
}